#include <cstdio>
#include <cstdlib>
#include <vector>
#include "utils.h"
#include <cuda_runtime.h>
#include <random>
#include <iostream>

#include "check_point/ck0/topk.cpp"

int main(int argc, char *argv[]) {
    cudaSetDevice(0);
    const int sm_per_block = 48*1024;
    std::vector<Case>  cases = {};

    #ifdef PERF
    float m_start = 100000;
    int loop = 1000;
    const int k = 20;
    
    for(int m_id = m_start; m_id <= 10*m_start; m_id+=m_start){
        cases.push_back({m_id, k, loop});
    }
    #endif

    #ifdef EVAL
    const int m = atoi(argv[1]), k = atoi(argv[2]), loop = atoi(argv[3]);
    const std::string src_gt_path = argv[4], dst_my_path = argv[5];
    cases.push_back({m, k, loop});
    #endif

    #ifdef PROFILE
    cases.push_back({262133, 20, 1});
    #endif

    for(size_t case_id = 0; case_id < cases.size(); ++case_id){
        #ifdef PERF
        if(case_id > 0){
            // cases[case_id].loop = max(1, int(1000.0f/cases[case_id-1].latency))*3;
            cases[case_id].loop = 1;
        }
        #endif

        auto &c = cases[case_id];
        const int m = c.m, k = c.k, loop = c.loop;
        cudaStream_t stream = NULL;
        CUDA_CHECK(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));
        float *d_src = nullptr;
        float *d_tmp_dst = nullptr;
        float *d_dst = nullptr;
        CUDA_CHECK(cudaMalloc(reinterpret_cast<void **>(&d_src), sizeof(float)*m));
        CUDA_CHECK(cudaMalloc(reinterpret_cast<void **>(&d_dst), sizeof(float)*k));
        
        const int raw_bt_sz = std::min(((sm_per_block/int(sizeof(float)))/k), 256);
        int bt_sz = 0;
        int start = 2;
        for(int i = 0; i < 8; ++i){
            if(start < raw_bt_sz){
                bt_sz = start;
            }else{
                break;
            }
            start <<=1;
        }
        if(bt_sz < 2){
            printf("### m %d k %d bt_sz %d invalid argument !!!\n", m, k, bt_sz);
            continue;
        }
        printf("bt_sz %d \n", bt_sz);
        
        int nk = DIV_UP(m, bt_sz*k);
        std::vector<int> its;
        its.push_back(m);
        if(nk > 1){
            CUDA_CHECK(cudaMalloc(reinterpret_cast<void **>(&d_tmp_dst), sizeof(float)*nk*k));
            for(int i = nk; i > 1; i = DIV_UP(i, bt_sz)){
                its.push_back(i*k);
            }
        }
        its.push_back(k);

        #ifdef EVAL
        loadbin(d_src, m*sizeof(float), src_gt_path);
        #endif
        {
            int start = now();
            for(int i = 0; i < loop; ++i){
                if(its.size() == 2){
                    topk(its[0], k, d_src, d_dst, bt_sz, stream);
                }else{
                    for(int i = 0; i < its.size(); ++i){
                        if(its[i] > k && its[i+1] != k){
                            if(i == 0){
                                topk(its[i], k, d_src, d_tmp_dst, bt_sz, stream);
                            }else{
                                topk(its[i], k, d_tmp_dst, d_tmp_dst, bt_sz, stream);
                            }
                        }else if(its[i] > k && its[i+1] == k){
                            topk(its[i], k, d_tmp_dst, d_dst, bt_sz, stream);
                        }
                    }
                }
            }
            cudaStreamSynchronize(stream);
            int end = now();
            c.latency =float(end-start)/float(loop);
            printf("m %d k %d loop %d latency %0.2f ms  \n",c.m, c.k, c.loop, c.latency);
            #ifdef EVAL
            dumpbin(d_dst, k*sizeof(float), dst_my_path);
            #endif
        }
        CUDA_CHECK(cudaFree(d_src));
        CUDA_CHECK(cudaFree(d_dst));
        if(!d_tmp_dst) CUDA_CHECK(cudaFree(d_tmp_dst));
        CUDA_CHECK(cudaStreamDestroy(stream));
    }
    return 0;
}